tree grammar ASTBuilder;

options {
  language = Java;
  tokenVocab = TBTParser;
  ASTLabelType = CommonTree;
}

@header {
  package typechecker.ast;
  
  import typechecker.ast.type.*;
  import typechecker.ast.term.*;
  
  import java.util.LinkedList;
  import java.util.Stack;
  import java.util.Map;
}

@members {
  private Context<TermVariable> vars = new Context<TermVariable>();
  {
    vars.push(new TermVariable("unit")); 
    vars.push(new TermVariable("bind"));  
  }
  private Context<TypeVariable> tvars = new Context<TypeVariable>(){
    public TypeVariable lookup(String s){
      TypeVariable tc = super.lookup(s);
      if(tc instanceof TypeConstructor){
			  /**
			    * Type constructors can be identified by name. 
			    * Every occurence has a different stage, so we must have different instances.
			    */
        return new TypeConstructor(s);
      }
      return tc;
    }
  };
  public void p(Object o){
    System.out.println(o);
  }
}

prog returns [Expression exp, List<InductiveDatatype> rdts]
@init{
  List<TermVariable> lvars    = new LinkedList<TermVariable>();
  List<TypeExpression> types  = new LinkedList<TypeExpression>();
  List<Expression> exps       = new LinkedList<Expression>();  
  List<InductiveDatatype> dts = new LinkedList<InductiveDatatype>();
} 
    : ^(PROG p=prolog[dts] ^(TL_BINDINGS (tl_bind[lvars, types, exps])*)  ^(TERM term) )
      { 
        /* Desugaring of top-level bindings into lambda abstraction-application */
        Expression texp = $term.exp;
        for(int i=lvars.size()-1; i >= 0; i--){
          texp = new LambdaAbstraction(lvars.get(i), types.get(i), texp);
          texp = new LambdaApplication(texp, exps.get(i));
        }
        $exp = texp;
        $rdts = dts;
        /* Clean up the top level bindings and data constructors */
        vars.clear();
        tvars.clear();
      }
    ;

tl_bind [List<TermVariable> lvars , List<TypeExpression> types, List<Expression> exps]
    : ^(TL_BINDING push_ID type term)        // XXX pushID
      { 
        lvars.add(vars.lookup($push_ID.text)); 
        exps.add($term.exp); 
        types.add($type.texp); 
      }
    ;

/* Gather information about defined datatypes */  
prolog [List<InductiveDatatype> dts]
    : ^(PROLOG (datatype[dts])*);    
    
datatype [List<InductiveDatatype> dts] 
@init{
  List<TypeVariable> params = new LinkedList<TypeVariable>();
  InductiveDatatype dt;
}
    : ^(DATATYPE  name=push_CTID 
                  args[params] { dt = new InductiveDatatype(new TypeConstructor($name.text), params); }
			            tconss[dt]
			)
			{ 
			 /* Add type parameters as parameters to constructors */
			 for( int i=0; i < params.size(); i++){
			   TypeVariable v = tvars.pop(); 
			   for( Map.Entry<DataConstructor, TypeExpression> kv : dt.entrySet()){
			     TypeExpression crt = kv.getValue();
			     kv.setValue(new ForallType(v, crt));
			   }
			 }
       dts.add(dt); 
			}
    ;   
args [ List<TypeVariable> params ] // arguments of the type constructor
    : ^(TARGS (arg=push_TID { params.add(tvars.lookup($arg.text));} )* ) ;  
tconss [InductiveDatatype dt] 
    : ^(TCONSTRUCTORS ( tcons[dt])*) ;
tcons [InductiveDatatype dt]
    : ^(TCONS push_CID type) { dt.put((DataConstructor)vars.lookup($push_CID.text), $type.texp);} ;

/* Generate AST for the lambda terms */
term  returns [Expression exp]
    : ^(CASE ty=type t=term { $exp = new CaseAnalisys($t.exp,$ty.texp); } branches[((CaseAnalisys)$exp)])     
    | ^(LETREC type push_ID t=term)    
          { $exp = new Letrec($type.texp, vars.lookup($push_ID.text), $t.exp); vars.pop(); }
    | ^(TYPE_ABS push_TID t=term)      
          { $exp = new TypeAbstraction(tvars.lookup($push_TID.text),$t.exp ); tvars.pop(); }
    | ^(TYPE_APP t=term ty=type)  
          { $exp = new TypeApplication($t.exp, $ty.texp); }
    | ^(LAM_ABS push_ID ty=type b=term)
          { $exp = new LambdaAbstraction(vars.lookup($push_ID.text), $ty.texp, $b.exp); vars.pop(); }
    | ^(LAM_APP t1=term t2=term)  
          { $exp = new LambdaApplication($t1.exp, $t2.exp); }  
    | ^(VAR ID)                   
          { $exp = vars.lookup($ID.text); }
    ;

branches [CaseAnalisys ca ] : ^(BRANCHES branch[ca]*) ;
branch   [CaseAnalisys ca ] 
    : ^(BRANCH ID t=term) // FIXME Lookup
      { ca.put(vars.lookup($ID.text), $t.exp); }
    ;

/* Initially in the AST there are only raw and starred types. Stages are added as they are inferred */
type returns [TypeExpression texp ]
    : ^(ARROW t1=type t2=type)  
          { $texp = new FunctionType($t1.texp, $t2.texp); }
    | ^(FOR_ALL push_TID t=type)      
          { $texp = new ForallType(tvars.lookup($push_TID.text), $t.texp); tvars.pop(); }
    | ^(TVAR TID)               
          { $texp = tvars.lookup($TID.text); }
    | ^(TVAR ^(STARED_TYPE TID))
          { TypeConstructor tt = (TypeConstructor) tvars.lookup($TID.text); tt.setStarred(); $texp = tt; }
    | tc=tcons_app
          { $texp = $tc.texp; } 
    ;
   
    
/**
  * We do not allow partial application of type constructors (kind checking) 
  */
tcons_app returns [TypeExpression texp ]
@init{
  List<TypeExpression> params = new LinkedList<TypeExpression>();
}
    // ty can be only be (stared) type variable 
    : ^(TCONS_APP ty=type (t=type {params.add($t.texp);} )*)    
          { $texp = new TypeConstructorApplication((TypeConstructor) $ty.texp, params); }
    ;
           
/* Match an identifier and push it into the appropriate context */
push_ID   returns [String text]  : ID  { vars.push(new TermVariable($ID.text));  $text=$ID.text; };
push_CID  returns [String text]  : ID  { vars.push((TermVariable)new DataConstructor($ID.text));  $text=$ID.text; };
push_TID  returns [String text]  : TID {tvars.push(new TypeVariable($TID.text)); $text=$TID.text;}; 
push_CTID returns [String text]  : TID {tvars.push((TypeVariable)new TypeConstructor($TID.text)); $text=$TID.text;}; 
